from PIL import Image
from torchvision import transforms

from config_manager import ImageAugmentConfig,TextAugmentConfig


class DataProcessor:
    """
    动态数据增强处理类。
    根据配置自动构建transform管线，可直接用于PyTorch的Dataset。
    """

    def __init__(self, config: ImageAugmentConfig|TextAugmentConfig,data_modal:str):
        self.config = config
        # 构建transform序列
        self.data_modal = data_modal
        self.transform = self._build_transform()

    def _build_transform(self):
        if self.data_modal == "image":
            return self._build_image_transform()
        elif self.data_modal == "text":
            return self._build_text_transform()
        else:
            return None

    def _build_image_transform(self):
        """
        根据配置项动态创建torchvision.transforms管线。
        每个操作可单独启用或禁用。
        """

        augment_list = []
        if self.config is not None and self.config.use_augmentation:

            # ============================================================
            # 1️⃣ 几何增强操作
            # 随机缩放
            if self.config.random_resize_crop:
                augment_list.append(transforms.RandomResizedCrop(224))

            # 裁剪（优先级高于普通随机裁剪）
            elif self.config.random_crop:
                augment_list.append(transforms.RandomCrop(32, padding=4))

            # 随机翻转
            if self.config.random_flip:
                augment_list.append(transforms.RandomHorizontalFlip(p=self.config.augment_prob))

            # 随机旋转
            if self.config.random_rotation:
                augment_list.append(transforms.RandomRotation(15))

            # 仿射变换
            if self.config.random_affine:
                augment_list.append(transforms.RandomAffine(degrees=10, translate=(0.1, 0.1)))

            # 透视变换
            if self.config.random_perspective:
                augment_list.append(
                    transforms.RandomPerspective(distortion_scale=0.3, p=self.config.augment_prob)
                )

            # ============================================================
            # 2️⃣ 颜色与光照增强 随机调整亮度、对比度、饱和度、色调（常用且有效的颜色增强）
            # brightness/contrast/saturation/hue 控制调整幅度（0.4/0.4/0.4/0.1 是经典配置）
            # - 亮度：模拟光照强弱变化
            # - 对比度：模拟曝光差异
            # - 饱和度：模拟相机色彩差异
            # - 色调：改变色彩倾向，轻微随机偏移
            if self.config.color_jitter:
                augment_list.append(
                    transforms.ColorJitter(
                        brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1
                    )
                )

            # 随机将图像转为灰度，模拟颜色丢失或低饱和度场景
            # p 控制概率，一般 0.1～0.3
            if self.config.grayscale:
                augment_list.append(transforms.RandomGrayscale(p=self.config.augment_prob))

            # 高斯模糊，模拟相机对焦不准、运动模糊等情况
            # kernel_size=3 适用于 CIFAR10（32×32 小图像）
            if self.config.gaussian_blur:
                augment_list.append(transforms.GaussianBlur(kernel_size=3))

            # 直方图均衡化（Equalize）：增强对比度，使暗部/亮部细节更均匀
            # 在弱光或高对比图像上有帮助
            if self.config.equalize:
                augment_list.append(transforms.RandomEqualize(p=self.config.augment_prob))

            # 太阳化（Solarize）：反转高亮部分像素（>threshold）
            # 模拟光照过强或曝光过度，常见于 AutoAugment / RandAugment 策略
            if self.config.solarize:
                augment_list.append(transforms.RandomSolarize(threshold=128, p=self.config.augment_prob))

            # 色度量化（Posterize）：降低图像的比特深度，使图像颜色块化
            # 这会迫使模型学习对颜色不敏感的特征
            if self.config.posterize:
                augment_list.append(transforms.RandomPosterize(bits=4, p=self.config.augment_prob))

            # ============================================================
            # 3️⃣ 随机擦除（仅作用于Tensor）
            # 是否将PIL或numpy转为Tensor（范围[0,1]）
            if self.config.to_tensor:
                augment_list.append(transforms.ToTensor())
                if self.config.random_erasing:
                    augment_list.append(transforms.RandomErasing(p=self.config.augment_prob))

            # 是否执行归一化，将像素范围从[0,1]映射为[-1,1]
            if self.config.normalize:
                augment_list.append(transforms.Normalize(self.config.mean, self.config.std))

        # 最终组合为Compose对象
        return transforms.Compose(augment_list)

    def _build_text_transform(self):
        pass

    def __call__(self, img):
        """
        对单张图像执行预处理与增强。
        支持PIL图像或numpy数组输入。
        """
        if not isinstance(img, Image.Image):
            img = Image.fromarray(img)  # numpy 转 PIL
        return self.transform(img)


